import numpy as np

from manim_sandbox.utils.mobjects.Trajectory import Trajectory, TrajectoryOnTime, TrajectoryOnLength
from manimlib import *
from manim_sandbox.utils.scenes.scene360 import *


class Bullet(Triangle):
    CONFIG = {
        "fill_opacity": 1,
        "stroke_width": 0,
        "length": DEFAULT_ARROW_TIP_LENGTH,
        "start_angle": PI,
        "aspect": 1.5,
        "angle": 0,
        "enter": 0
    }

    def __init__(self, location=ORIGIN, **kwargs):
        super().__init__(**kwargs)
        digest_config(self, kwargs)
        self.set_height(self.length, stretch=True)
        self.set_width(self.length * self.aspect, stretch=True)
        self.data["points"][4] += np.array([self.length, 0, 0])
        self.scale(0.5)
        self.move_to(location)

    def get_angle(self):
        return self.angle
        # return angle_of_vector(self.get_vector())

    def get_vector(self):
        return self.point_from_proportion(0.5) - self.get_start()

    def rotate(self, angle, axis=OUT, **kwargs):
        super().rotate(angle, axis, **kwargs)
        self.angle = (self.angle + angle) % TAU
        return self


class WaveAndParticle(Scene):
    def construct(self):
        vqueue1 = VGroup()
        vqueue2 = VGroup()
        color_list = [RED, ORANGE, GOLD, YELLOW, GREEN, TEAL, BLUE, PURPLE]
        color_list.reverse()
        self.an = 0

        def update_bullet1(obj, dt):
            if self.time < 40:
                obj.add(Bullet().scale(0.7).rotate(self.time ** 1.4 * 3 * PI))
            obj.set_color_by_gradient(*color_list)
            for k in obj:
                k.shift(dt * 2 * np.array([
                    np.cos(k.get_angle()),
                    np.sin(k.get_angle()),
                    0
                ]))
                if abs(get_norm(k.get_center() - ORIGIN)) > 7.5:
                    obj.remove(k)

        def update_bullet2(obj, dt):
            if self.time < 40:
                obj.add(Bullet().scale(0.7).rotate(self.time ** 1.4 * 3 * PI + PI))
            obj.set_color_by_gradient(*color_list)
            for k in obj:
                k.shift(dt * 2 * np.array([
                    np.cos(k.get_angle()),
                    np.sin(k.get_angle()),
                    0
                ]))
                if abs(get_norm(k.get_center() - ORIGIN)) > 7.5:
                    obj.remove(k)

        vqueue1.add_updater(update_bullet1)
        vqueue2.add_updater(update_bullet2)
        self.add(vqueue1, vqueue2)
        self.wait(45)


class WaveAndParticleWithSpirals(Scene):
    def construct(self):
        vqueue1 = VGroup()
        vqueue2 = VGroup()
        # vqueue3 = VGroup()
        self.count = 0

        def one_trail(start_angle):
            def update(obj, dt):
                if self.count < 60:
                    self.count += 1
                    return
                obj.add(Bullet().scale(0.7).rotate(self.time ** 2 * 0.5 + start_angle))
                obj.set_color_by_gradient(RED, PURPLE, BLUE, TEAL, GREEN, YELLOW, GOLD, ORANGE, RED)

                for k in obj:
                    k.shift(dt * 2 * np.array([
                        np.cos(k.get_angle()),
                        np.sin(k.get_angle()),
                        0
                    ]))
                    k.rotate(-2 * DEGREES)
                    if abs(get_norm(k.get_center() - ORIGIN)) > 7:
                        obj.remove(k)
                    k.enter += dt
                    if k.enter >= 6:
                        obj.remove(k)
                self.count += 1

            return update

        # self.wait(0.5)
        vqueue1.add_updater(one_trail(0))
        vqueue2.add_updater(one_trail(TAU / 2))
        # vqueue3.add_updater(one_trail(TAU / 3 * 2))
        self.add(vqueue1)
        self.add(vqueue2)
        # self.add(vqueue3)
        # self.wait(10)


class BallBullet(Sphere):
    CONFIG = {
        "resolution": (10, 10),
    }

    def __init__(self, position: np.ndarray = ORIGIN, **kwargs):
        super().__init__(**kwargs)
        self.move_to(position)
        self.set_width(0.25)
        self.__theta = 0
        self.__gamma = 0
        self.__lifetime = 0

    def get_theta(self):
        return self.__theta

    def get_gamma(self):
        return self.__gamma

    def set_theta(self, theta):
        self.__theta = theta
        return self

    def incre_theta(self, d_theta):
        self.__theta += d_theta
        return self

    def set_gamma(self, gamma):
        self.__gamma = gamma
        return self

    def incre_gamma(self, d_gamma):
        self.__gamma += d_gamma
        return self

    def life_time(self):
        return self.__lifetime

    def incre_lifetime(self, d_time):
        self.__lifetime += d_time
        return self


class BallBullet2(Sphere):
    CONFIG = {
        "resolution": (10, 10),
    }

    def __init__(self, position: np.ndarray = ORIGIN, **kwargs):
        super().__init__(**kwargs)
        self.move_to(position)
        self.set_width(0.25)
        self.vx, self.vy, self.vz = 0, 0, 0


class Danmaku360(Scene):
    def construct(self):
        g1 = Group()
        g2 = Group()
        g3 = Group()
        self.count = 0

        def update_one_tail(start_angle=0):
            def update(obj: Group, dt):
                if self.count < 9:
                    self.count += 1
                    return
                obj.add(BallBullet(IN * 6).set_theta(self.time ** 2 / 2 + start_angle))
                obj.set_color_by_gradient(RED, PURPLE, BLUE, TEAL, GREEN, YELLOW, GOLD, ORANGE, RED)
                for k in obj:
                    k.shift(np.array([
                        np.cos(k.get_theta()),
                        np.sin(k.get_theta()),
                        np.sin(k.get_gamma())
                    ]) * 4 * dt)

                    k.incre_gamma(dt * PI / 6)
                    k.incre_theta(-1 * DEGREES)

                    k.incre_lifetime(dt)
                    if k.life_time() >= 6:
                        obj.remove(k)

            return update

        g1.add_updater(update_one_tail(0))
        g2.add_updater(update_one_tail(TAU / 3))
        g3.add_updater(update_one_tail(TAU / 3 * 2))
        self.add(g1, g2, g3)
        # self.add(g1)
        self.wait(30)


class LorenzScene(Scene):
    def construct(self):
        g = Group()
        self.add(ThreeDAxes())
        # self.camera.frame.shift(OUT*3)
        self.t = 0
        a, b, c = 0.2, 0.2, 5.6
        pos = np.array([3, 2, 1])

        def add_bullet(obj: Group, dt):
            self.t += 1
            if self.t % 6 == 0:
                obj.add(BallBullet2(pos))

        def update_one_trace(obj: Group, dt):
            obj.set_color_by_gradient(RED, PURPLE, BLUE, TEAL, GREEN, YELLOW, GOLD, ORANGE, RED)
            for k in obj:
                k.vx = (-k.get_center()[1] - k.get_center()[2])
                k.vy = (k.get_center()[0] + a * k.get_center()[1])
                k.vz = (b + k.get_center()[2] * (k.get_center()[0] - c))
                k.move_to(
                    k.get_center() + np.array([k.vx, k.vy, k.vz]) * dt
                )

        g.add_updater(add_bullet)
        g.add_updater(update_one_trace)
        self.add(g)
        self.wait(50)
        g.remove_updater(add_bullet)
        self.wait(50)


class VBullet(Bullet):
    def __init__(self, location=ORIGIN, **kwargs):
        super().__init__(location, **kwargs)
        self.data["velocity"] = ORIGIN
        self.data["acceleration"] = ORIGIN

    def get_velocity(self):
        return self.data["velocity"]

    def set_velocity(self, velocity: np.ndarray):
        self.data["velocity"] = velocity
        self.set_angle(angle_of_vector(velocity))
        return self

    def incre_velocity(self, velocity: np.ndarray):
        self.set_velocity(self.data["velocity"] + velocity)
        return self

    def get_acceleration(self):
        return self.data["acceleration"]

    def set_acceleration(self, acceleration: np.ndarray):
        self.data["acceleration"] = acceleration
        return self

    def step_velocity(self, dt):
        self.incre_velocity(self.get_acceleration() * dt)
        return self

    def step(self, dt):
        self.step_velocity(dt)
        self.shift(self.get_velocity() * dt)
        return self

    def set_angle(self, angle, about_point=None):
        if about_point is None:
            about_point = self.get_start()
        self.rotate(
            angle - self.get_angle(),
            about_point=about_point,
        )
        return self


def polar2xyz(rho, theta):
    return rho * np.array([np.cos(theta), np.sin(theta), 0])


INF = 2 ** 8


class ElectricFieldScene(Scene):
    def construct(self):
        k = 1
        q1 = 1
        q2 = -12
        q0 = 1
        m = 1
        f1 = TrueDot(LEFT * 1).make_3d()
        f2 = TrueDot(RIGHT * 1).make_3d()
        self.add(f1, f2)

        # self.add(f1)

        def update_state(obj: VBullet, dt):
            force1 = clip(- k * q1 * q0 / norm_squared(f1.get_center() - obj.get_center()), -INF, INF)
            force2 = clip(- k * q2 * q0 / norm_squared(f2.get_center() - obj.get_center()), -INF, INF)

            angle1 = angle_of_vector(f1.get_center() - obj.get_center())
            angle2 = angle_of_vector(f2.get_center() - obj.get_center())

            vf1 = polar2xyz(force1, angle1)
            vf2 = polar2xyz(force2, angle2)

            obj.set_acceleration((vf1 + vf2) / m)
            # obj.set_acceleration(vf1 / m)
            obj.step(dt)

        bullets = VGroup(*[
            VBullet().set_velocity(RIGHT * 2).scale(0.5) for i in range(64)
        ]).arrange(DOWN, buff=0.08).to_edge(LEFT, buff=3).set_color_by_gradient(RED, YELLOW, GREEN, TEAL)

        for i in range(len(bullets)):
            bullets[i].add_updater(update_state)

        paths = VGroup(*[
            TracedPath(bullets[i].get_center, stroke_width=[0, 1], stroke_opacity=[0, 0.5], time_per_anchor=1 / 5)
            for i in range(len(bullets))
        ])

        self.add(*bullets)
        self.add(*paths)

        # self.wait(10)
